import pandas as pd
from MyTT import *
import akshare as ak
import numpy as np
import os
from datetime import datetime
from sqlalchemy import create_engine, distinct, or_, and_
import sqlite3
import backtrader as bt
import pymssql
from urllib.parse import quote_plus as urlquote
from configparser import ConfigParser

pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)

conf = ConfigParser()
conf.read('env.ini')

sqlserver = ('sqlserver_out', 'sqlserver_in')[0]

host = conf.get(sqlserver, 'host')
port = conf.get(sqlserver, 'port')
user = conf.get(sqlserver, 'user')
password = conf.get(sqlserver, 'password')
schema = conf.get(sqlserver, 'schema')
db_name = conf.get(sqlserver, 'db')


class TestStrategy(bt.Strategy):
    params = (('maperiod', 10), ('lb_v', 2.2), ('pre_ratio', 0.0048),
              ('ratio', 0.1), ('trail_ratio', 0.035), ('printlog', False),)

    def __init__(self):
        # 初始化交易指令
        self.order = None
        self.buy_list = []
        # self.buy_stock = get_stocks()
        # self.trade_dates = self.buy_stock['trade_date'].unique().tolist()
        self.order_list = []  # 记录以往订单，方便调仓日对未完成订单做处理
        self.buy_price = 1e4  # 每个证券买入金额，不满一手成交一手
        self.daily_pnl = dict()
        self.trade_info = list()
        self.buy_info = dict()  # 证券买入信息，包括买入时间和买入成本，例如{‘000001’: {'date': 20230614, 'cost': 14.23}}
        self.lb = dict()
        self.lb_count = dict()
        self.lb_fit = dict()
        self.period_high = dict()
        self.code_high = dict()
        self.price_high = dict()
        self.base_line = dict()  # 每个证券止损线
        self.cache_log = list()

        for data in self.datas:
            data_name = data._name
            if data_name[-2:] == 'pl':
                self.period_high[data] = bt.indicators.Highest(data.high, period=28, subplot=False)
            else:
                self.lb[data] = bt.DivByZero(data.volume, bt.ind.SMA(data.volume, period=self.p.maperiod))

    def downcast(self, price):
        res = self.buy_price / price // 100 * 100
        res = res if res > 0 else 100
        return res

    def next(self):
        if self.order:  # 查看是否有指令执行，如果有则不执行这bar
            return

        # for code in set(self.getdatanames()) - set(self.buy_list):
        for i, data in enumerate(self.datas):
            log_str = ""
            data_name = data._name
            code = data_name[:6]
            if data_name[-2:] == 'pl':
                continue
            else:
                data_pl = self.datas[i + 1]
                beg_time = data_pl.lines.datetime.datetime(0)
                if beg_time.hour == 15 and beg_time.minute == 0:
                    self.code_high[data] = self.period_high[data_pl][0]
                if (data not in self.lb_count) or (beg_time.hour == 9 and beg_time.minute == 30) or (
                        beg_time.hour == 13 and beg_time.minute == 1):
                    self.lb_count[data] = 1
                    self.lb_fit[data] = False
                    if self.lb[data] > self.p.lb_v:
                        self.lb_fit[data] = True
                elif self.lb[data] > self.p.lb_v:
                    self.lb_fit[data] = True
                    self.lb_count[data] = 1
                else:
                    self.lb_count[data] += 1

                if self.lb_count[data] > 5:
                    self.lb_fit[data] = False

            dt = data.datetime.date(0)  # 获取当前的回测时间点
            dt = int(dt.strftime('%Y%m%d'))

            # self.log(
            #     f'code={code}->datetime {data.lines.datetime.datetime(0)}, close {data.close[0]}, 资产：{self.broker.getvalue()} 元')

            cur_datetime = data.lines.datetime.datetime(0)
            if cur_datetime.hour == 15:
                if dt not in self.daily_pnl:
                    # item = [dt, self.broker.getvalue() - init_cash]
                    self.daily_pnl[dt] = round(self.broker.getvalue() - init_cash, 2)
                if code in self.buy_info:
                    self.buy_info.clear()

            # print(f"-------------{code} next 的第{count_info[code] + 1}次循环 --------------")
            # print(f"今日：datetime {data.lines.datetime.date(0)} close {data.lines.close[0]} high {data.lines.high[0]}")
            # print(
            #     f"code={code}->当前时点（今日）：datetime {data.lines.datetime.datetime(0)} close {data.lines.close[0]} high={data.lines.high[0]} open={data.lines.open[0]}->lb={self.lb[data][0]}")
            # log_str += f"-------------{code} next 的第{count_info[code] + 1}次循环 --------------\n"
            # log_str += f"code={code}->今日：datetime {data.lines.datetime.datetime(0)} close {data.lines.close[0]} high {data.lines.high[0]} lb {self.lb[data][0]} lb_fit {self.lb_fit[data]}\n"
            # if data in self.code_high:
            #     log_str += f"code={code}->当前时点（今日）：datetime {data.lines.datetime.datetime(0)} close {data.lines.close[0]}->high={data.lines.high[0]}->ever_high={self.code_high[data]}->lb={self.lb[data][0]}->lb_fit={self.lb_fit[data]}\n"
            #     log_str += f'time={cur_datetime}->data in code_high={data in self.code_high} data in self.lb_fit = {data in self.lb_fit}->f1={data.high > self.code_high[data] * (1 + pre_ratio)}->f2={self.lb_fit[data]}'
            #     self.cache_log.append(log_str)
            # if code in buy_info:
            #     print(f"code={code}->time={cur_datetime}->low={data.low[0]}->line={buy_info[code]['cost'] * (1 - ratio)}->close={data.close[0]}->hl={buy_info[code]['cost'] * (1 + ratio / 2)}->trail_price={buy_info[code]['cost'] * (1 + trail_ratio)}")
            #     print(f'buy_info={buy_info}')

            if code in self.price_high and data.high[0] > self.price_high[code]:
                self.price_high[code] = data.high[0]
                new_line = (1 - self.p.ratio / 2) * data.high[0]
                self.base_line[code] = new_line
                # if new_line > code_high[code]:
                #    code_line[code] = new_line
                # print(
                #     f"INLINE->code={code}->time={cur_datetime}->low={data.low[0]}->line={buy_info[code]['cost'] * (1 - ratio)}->close={data.close[0]}->hl={buy_info[code]['cost'] * (1 + ratio / 2)}->trail_price={buy_info[code]['cost'] * (1 + trail_ratio)}")
                # continue

            pos = self.getposition(data).size
            # print(
            #     f"今日：datetime {data.lines.datetime.datetime(0)} close {data.lines.close[0]} high {data.lines.high[0]}")
            # print(f"今日：h {cur_datetime.hour} m {cur_datetime.minute} pos {pos}")

            if (not pos) and (dt in code_info[code]) and (code not in self.buy_info) and (
                    (cur_datetime.hour != 14) or (cur_datetime.hour == 14 and cur_datetime.minute <= 45)):
                # if data in self.code_high and data in self.lb_fit:
                # print(f'time={cur_datetime}->data in code_high={data in self.code_high} data in self.lb_fit={data in self.lb_fit}->f1={data.high > self.code_high[data] * (1 + pre_ratio)}->f2={self.lb_fit[data]}')
                if cur_datetime.hour != 15 and data in self.code_high and data in self.lb_fit and data.high > \
                        self.code_high[data] * (
                        1 + self.p.pre_ratio) and self.lb_fit[data]:
                    self.order = self.buy(data=data, size=100, name=code)
                    self.buy_info[code] = {'date': dt}
                    self.price_high[code] = data.high[0]
                    self.base_line[code] = (1 - self.p.ratio / 2) * data.high[0]
                    # print(
                    #     f'code: {code} datetime {data.lines.datetime.datetime(0)} BUY,BUY,BUY!!! high={data.high[0]}->pro_high={code_info[code][dt]}->lb={self.lb[data][0]}')
                    # log_str += f'BUY,BUY,BUY!!! high={data.high[0]}->pro_high={code_info[code][dt]}->lb={self.lb[data][0]}\n'

            elif code in self.buy_info and data.low < self.base_line[code]:
                self.order = self.close(data=data, name=code)
            elif cur_datetime.hour == 14 and cur_datetime.minute > 50:
                self.order = self.close(data=data, name=code)
            # print(
            #     f'code: {code} datetime {data.lines.datetime.datetime(0)} SELL,SELL,SELL!!! base={code_line[code]}->low={data.low[0]}->lb={self.lb[data][0]}')
            # log_str += f'SELL,SELL,SELL!!! base={code_line[code]}->low={data.low[0]}\n'

            # 在 next() 中调用 len(self.data0)，返回的是当前已处理（已回测）的数据长度，会随着回测的推进动态增长
            # print(f"已处理的数据点：{len(data)}")
            # log_str += f"已处理的数据点：{len(data)}\n"
            # # buflen() 返回整条线的总长度，固定不变；
            # print(f"line的总长度：{data.buflen()}")
            # log_str += f"line的总长度：{data.buflen()}\n"

    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            # print(f'{dt.isoformat()},{txt}')
            print(txt)

    def notify_order(self, order):
        # 未被处理的订单
        if order.status in [order.Submitted, order.Accepted]:
            return
        # 已经处理的订单
        if order.status in [order.Completed, order.Canceled, order.Margin]:
            code = order.data._name[:6]
            if order.isbuy():
                self.log(
                    f'code={code} BUY CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log(
                    'BUY EXECUTED, ref:%.0f，Price: %.2f, Cost: %.2f, Comm %.2f, Size: %.2f, Stock: %s' %
                    (order.ref,  # 订单编号
                     order.executed.price,  # 成交价
                     order.executed.value,  # 成交额
                     order.executed.comm,  # 佣金
                     order.executed.size,  # 成交量
                     code))  # 股票名称
                trade_dt = int(bt.num2date(order.created.dt).strftime('%Y%m%d%H%M%S'))
                trade_item = [trade_dt, code, 'BUY', round(order.executed.price, 2), order.executed.size,
                              round(order.executed.comm, 2)]
                self.buy_info[code]['cost'] = order.executed.price
                self.trade_info.append(trade_item)
            else:  # Sell
                self.log(
                    f'code={code} SELL CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log('SELL EXECUTED, ref:%.0f, Price: %.2f, Cost: %.2f, Comm %.2f, Size: %.2f, Stock: %s' %
                         (order.ref,
                          order.executed.price,
                          order.executed.value,
                          order.executed.comm,
                          order.executed.size,
                          code))
                trade_dt = int(bt.num2date(order.created.dt).strftime('%Y%m%d%H%M%S'))
                trade_item = [trade_dt, code, 'SELL', round(order.executed.price, 2),
                              order.executed.size, round(order.executed.comm, 2)]
                self.trade_info.append(trade_item)
            self.bar_executed = len(self)

        # Write down: no pending order
        self.order = None

    def stop(self):
        print('Strategy Finish!!!')

        with open(r'C:\Users\AndrewX\Desktop\cache_log.txt', 'w') as f:
            res = '\n'.join(self.cache_log)
            f.write(res)

        daily_pnl_df = pd.DataFrame(self.daily_pnl.items(), columns=['DATE_T', 'PNL'])
        append_df = pd.DataFrame([[101, round((self.broker.getvalue() / init_cash - 1) * 100, 2)]],
                                 columns=['DATE_T', 'PNL'])
        daily_pnl_df = daily_pnl_df.append(append_df, ignore_index=True)
        trade_df = pd.DataFrame(self.trade_info, columns=['DATE_T', 'STK_CODE', 'ACTION', 'PRICE', 'AMOUNT', 'COMM'])
        pnl_table = 'BREAKUPKZZ_DAILY_PNL'
        trade_table = 'BREAKUPKZZ_TRADE_INFO'
        daily_pnl_df = daily_pnl_df.loc[(daily_pnl_df['DATE_T'] >= 20230615) | (daily_pnl_df['DATE_T'] == 101)]
        insert_db(daily_pnl_df, pnl_table)
        insert_db(trade_df, trade_table)
        print('insert table ok')


def get_stocks_info():
    sql = f"SELECT * FROM KZZ_POOL WHERE trade_date>=20230616"
    r = exec_sql(sql)
    df = pd.DataFrame(r, columns=['trade_date', 'code', 'value'])
    df.sort_values('trade_date', inplace=True, ignore_index=True)

    return df


def get_data(stocks, freq='day'):
    if freq == '1min':
        table_name = 'KZZ_MIN_DATA'
    elif freq == '15min':
        table_name = 'KZZ_15MIN_DATA'
    else:
        table_name = 'KZZ_15MIN_DATA'

    test_code = '123197'
    stocks_str = ','.join(f"'{tok}'" for tok in stocks)
    sql = f"SELECT * FROM {table_name} WHERE code in ({stocks_str})"
    r = exec_sql(sql)
    df = pd.DataFrame(r, columns=['code', 'datetime', 'close', 'open', 'low', 'high', 'volume', 'money'])
    df['datetime'] = df['datetime'].map(time_map)
    sql2 = f"SELECT COUNT(1) FROM {table_name} WHERE code='{test_code}'"
    r2 = exec_sql(sql2)
    test_count = r2[0][0]
    res = dict()

    for code in stocks:
        sig_df = df.loc[df['code'] == code]
        print(f'code={code}->len={len(sig_df)}->test_count={test_count}')
        if len(sig_df) >= test_count or code == test_code:
            sig_df.sort_values('datetime', inplace=True, ignore_index=True)
            sig_df.set_index('datetime', inplace=True)
            res[code] = sig_df

    # print(f'res={res}')
    return res


def exec_sql(sql):
    conn = pymssql.connect(host=host, port=port, user=user, password=password, database=db_name)
    cursor = conn.cursor()
    cursor.execute(sql)
    r = cursor.fetchall()
    cursor.close()
    conn.close()
    return r


def time_map(x):
    t = datetime.strptime(str(x), '%Y%m%d%H%M')
    return t


def insert_db(df, table_name):
    yconnect = create_engine(f'mssql+pymssql://{user}:{urlquote(password)}@{host}:{port}/{db_name}?charset=utf8')
    pd.io.sql.to_sql(df, table_name, yconnect, schema=schema, if_exists='append', index=False)


code_info = dict()
init_cash = 1e8


def run_backtrade():
    cerebro = bt.Cerebro()
    stocks_info = get_stocks_info()
    stocks = stocks_info['code'].unique()

    data_1min = get_data(stocks, freq='1min')
    data_15min = get_data(stocks, freq='15min')

    for code in data_1min:
        data0 = data_1min[code]
        data1 = data_15min[code]
        datafeed0 = bt.feeds.PandasData(dataname=data0)
        cerebro.adddata(datafeed0, name=f'{code}_ps')
        datafeed1 = bt.feeds.PandasData(dataname=data1)
        cerebro.adddata(datafeed1, name=f'{code}_pl')

        df = stocks_info.query(f"code=='{code}'")
        if len(df) > 0:
            code_info[code] = df['trade_date'].unique()

        print(f'{code} feeds ok')

    ##   st_date = datetime.datetime(2023, 3, 1)
    ##ed_date = datetime.datetime(2023, 6, 5)

    # 初始资金 100,000,0
    cerebro.broker.setcash(init_cash)
    # 佣金，双边各 0.0003
    cerebro.broker.setcommission(commission=0.0003)
    # 滑点：双边各 0.0001
    cerebro.broker.set_slippage_perc(perc=0.0001)

    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='pnl')  # 返回收益率时序数据
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name='_AnnualReturn')  # 年化收益率
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='_SharpeRatio')  # 夏普比率
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='_DrawDown')  # 回撤

    cerebro.addstrategy(TestStrategy)

    result = cerebro.run()
    strat = result[0]
    # 返回日度收益率序列
    daily_return = pd.Series(strat.analyzers.pnl.get_analysis())
    # 打印评价指标
    print("--------------- TimeReturn -----------------")
    print(daily_return)
    print("--------------- AnnualReturn -----------------")
    print(strat.analyzers._AnnualReturn.get_analysis())
    print("--------------- SharpeRatio -----------------")
    print(strat.analyzers._SharpeRatio.get_analysis())
    print("--------------- DrawDown -----------------")
    print(strat.analyzers._DrawDown.get_analysis())
    # cerebro.plot()


if __name__ == '__main__':
    run_backtrade()
